import numpy as np
import matplotlib.pyplot as plt

class MPCController:
    def __init__(self, A, B, C, N, Q, R, u_min, u_max, reg=1e-5):
        self.A = A
        self.B = B
        self.C = C
        self.N = N
        self.Q = Q
        self.R = R
        self.u_min = u_min
        self.u_max = u_max
        self.reg = reg  # 正则化参数
        self.measurement_history = []
        self.setpoint_history = []

    def mpc_control(self, x, r):
        # n = self.A.shape[0]
        # m = self.B.shape[1]
        #
        # # 构建预测矩阵
        # P = np.zeros((self.N * n, n))
        # Qb = np.zeros((self.N * n, self.N * n))
        # Rb = np.zeros((self.N * m, self.N * m))
        # Pb = np.zeros((self.N * n, self.N * m))
        #
        # A_power = np.eye(n)  # 初始化为单位矩阵
        # for i in range(self.N):
        #     A_power = A_power @ self.A
        #     P[i * n:(i + 1) * n, :] = A_power
        #     Qb[i * n:(i + 1) * n, i * n:(i + 1) * n] = self.Q
        #
        #     B_power = np.eye(n)  # 初始化为单位矩阵
        #     for j in range(i + 1):
        #         if i - j - 1 >= 0:
        #             B_power = np.eye(n)  # 重置为单位矩阵
        #             for k in range(i - j):
        #                 B_power = B_power @ self.A
        #         else:
        #             B_power = np.eye(n)
        #         Pb[i * n:(i + 1) * n, j * m:(j + 1) * m] = B_power @ self.B
        #
        # for i in range(self.N):
        #     Rb[i * m:(i + 1) * m, i * m:(i + 1) * m] = self.R
        #
        # # 将参考轨迹展开成列向量
        # r = np.tile(r, (n, 1)).flatten()
        #
        # # 目标函数中的H矩阵和f向量
        # H = 2 * (Pb.T @ Qb @ Pb + Rb + self.reg * np.eye(self.N * m))
        #
        # f = 2 * (Pb.T @ Qb @ (P @ x0.flatten() - r))
        #
        # # 求解优化问题
        # u_opt = np.linalg.solve(H, -f)
        #
        # # 提取第一步的控制输入
        # u = u_opt[:m]
        # return u
        x = x.reshape(-1,1)
        r = r.reshape(-1,1)
        n_states = self.A.shape[0]
        n_controls = self.B.shape[1]

        # 构建优化问题
        H = np.zeros((self.N * n_controls, self.N * n_controls))
        F = np.zeros((self.N * n_controls, 1))

        for i in range(self.N):
            H[i * n_controls:(i + 1) * n_controls, i * n_controls:(i + 1) * n_controls] = self.R
            F[i * n_controls:(i + 1) * n_controls, 0] = 0

        G = np.zeros((self.N * n_states, self.N * n_controls))
        E = np.zeros((self.N * n_states, n_states))
        L = np.zeros((self.N * n_states, 1))

        for i in range(self.N):
            for j in range(i + 1):
                G[i * n_states:(i + 1) * n_states, j * n_controls:(j + 1) * n_controls] = np.linalg.matrix_power(self.A,
                                                                                                                 i - j) @ self.B
            E[i * n_states:(i + 1) * n_states, :] = np.linalg.matrix_power(self.A, i + 1)
            L[i * n_states:(i + 1) * n_states, :] = np.linalg.matrix_power(self.A, i + 1) @ x - r

        Q_bar = np.kron(np.eye(self.N), self.Q)
        H += G.T @ Q_bar @ G
        F += (G.T @ Q_bar @ L).reshape(-1, 1)

        # 计算最优控制输入
        u_opt = -np.linalg.inv(H) @ F

        # 取第一个控制输入
        u = u_opt[:n_controls]

        return u

    def update(self, setpoint, measurement):
        self.measurement_history.append(measurement)
        self.setpoint_history.append(setpoint)
        u = self.mpc_control(np.array(measurement), np.array(setpoint))
        # print(u)
        return u

    def plot_error(self, title='none'):
        self.setpoint_history = np.array(self.setpoint_history)
        self.measurement_history = np.array(self.measurement_history)
        for i in range(len(self.setpoint_history[0])):
            fig, ax = plt.subplots()
            ax.plot([i / 100 for i in range(len(self.setpoint_history[:, i]))], self.setpoint_history[:, i],
                    label='setpoint')
            ax.plot([i / 100 for i in range(len(self.measurement_history[:, i]))], self.measurement_history[:, i],
                    label='measurement')
            ax.set_xlabel('time')
            ax.set_ylabel('value')
            ax.set_title(title)
            ax.legend()

if __name__ == '__main__':
    A = np.array([[1, 0,0], [0, 1,0],[0,0,1]])
    B = np.array([[0.5,0.6,0.2], [1,0.1,0.4],[0.5,0.5,0.8]])
    C = np.eye(3)  # 输出矩阵
    Q = np.eye(3)*2
    R = np.eye(1)
    N = 10
    u_min = -1
    u_max = 1
    reg = 1e-5

    outer_controller = MPCController(A, B, C, N, Q, R, u_min, u_max, reg)
    inner_controller = MPCController(A, B, C, N, Q, R, u_min, u_max, reg)
    setpoint = np.array([[1], [1],[1]])
    # r = np.ones((1, N))  # 目标为1
    measurement = np.array([[0],[0],[0]])
    x = np.array([[0], [0],[0]])
    x_dot = np.array([[0], [0], [0]])
    for _ in range(10):
        u = outer_controller.update(setpoint, x)
        u_v = inner_controller.update(u,x_dot)
        x_dot = A @ x_dot + B @ u_v.reshape(-1, 1)
        x = A @ x + B @ x_dot.reshape(-1, 1)
        # x = A @ x + B @ u.reshape(-1, 1)
        y = C @ x
        print("控制输入 u:")
        print(u)
        print('状态 x:')
        print(x)
    inner_controller.plot_error()
    outer_controller.plot_error()
    plt.show()
